//! 查询抽象层
//! 
//! 定义了统一的查询接口和数据结构，支持：
//! - SQL查询
//! - GraphQL查询
//! - REST API查询
//! - Flight SQL查询

use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use arrow::record_batch::RecordBatch;

use crate::error::{KernelError, Result};

/// 查询类型
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum QueryType {
    /// SQL查询
    Sql,
    /// GraphQL查询
    GraphQL,
    /// REST API查询
    Rest,
    /// Flight SQL查询
    FlightSql,
    /// 脚本查询
    Script,
}

/// 查询请求
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Query {
    /// 查询ID
    pub id: String,
    /// 查询类型
    pub query_type: QueryType,
    /// 查询内容
    pub content: String,
    /// 查询参数
    pub parameters: HashMap<String, QueryParameter>,
    /// 查询选项
    pub options: QueryOptions,
    /// 查询上下文
    pub context: QueryContext,
}

/// 查询参数
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QueryParameter {
    String(String),
    Integer(i64),
    Float(f64),
    Boolean(bool),
    Null,
    Array(Vec<QueryParameter>),
    Object(HashMap<String, QueryParameter>),
}

/// 查询选项
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryOptions {
    /// 超时时间（毫秒）
    pub timeout_ms: Option<u64>,
    /// 最大结果行数
    pub max_rows: Option<usize>,
    /// 是否流式返回
    pub streaming: bool,
    /// 缓存策略
    pub cache_strategy: CacheStrategy,
    /// 执行模式
    pub execution_mode: ExecutionMode,
}

/// 缓存策略
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CacheStrategy {
    /// 不缓存
    None,
    /// 内存缓存
    Memory { ttl_seconds: u64 },
    /// 分布式缓存
    Distributed { ttl_seconds: u64 },
    /// 自动选择
    Auto,
}

/// 执行模式
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ExecutionMode {
    /// 同步执行
    Sync,
    /// 异步执行
    Async,
    /// 批量执行
    Batch,
    /// 流式执行
    Stream,
}

/// 查询上下文
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryContext {
    /// 用户ID
    pub user_id: Option<String>,
    /// 会话ID
    pub session_id: Option<String>,
    /// 请求ID
    pub request_id: String,
    /// 数据源ID
    pub datasource_id: Option<String>,
    /// 权限信息
    pub permissions: Vec<String>,
    /// 自定义属性
    pub attributes: HashMap<String, String>,
}

/// 查询结果
#[derive(Debug, Clone)]
pub struct QueryResult {
    /// 查询ID
    pub query_id: String,
    /// 执行状态
    pub status: QueryStatus,
    /// 结果数据
    pub data: QueryResultData,
    /// 执行统计
    pub stats: QueryStats,
    /// 错误信息
    pub error: Option<String>,
}

/// 查询状态
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QueryStatus {
    /// 等待执行
    Pending,
    /// 执行中
    Running,
    /// 执行成功
    Success,
    /// 执行失败
    Failed,
    /// 已取消
    Cancelled,
}

/// 查询结果数据
#[derive(Debug, Clone)]
pub enum QueryResultData {
    /// Arrow RecordBatch
    Arrow(Vec<RecordBatch>),
    /// JSON数据
    Json(serde_json::Value),
    /// 原始字符串
    Raw(String),
    /// 空结果
    Empty,
}

/// 查询执行统计
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryStats {
    /// 开始时间
    pub start_time: chrono::DateTime<chrono::Utc>,
    /// 结束时间
    pub end_time: Option<chrono::DateTime<chrono::Utc>>,
    /// 执行时间（毫秒）
    pub duration_ms: Option<u64>,
    /// 处理行数
    pub rows_processed: Option<usize>,
    /// 返回行数
    pub rows_returned: Option<usize>,
    /// 内存使用（字节）
    pub memory_used: Option<usize>,
    /// CPU时间（毫秒）
    pub cpu_time_ms: Option<u64>,
    /// 是否使用缓存
    pub cache_hit: bool,
}

impl Query {
    /// 创建SQL查询
    pub fn sql(content: String) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            query_type: QueryType::Sql,
            content,
            parameters: HashMap::new(),
            options: QueryOptions::default(),
            context: QueryContext::default(),
        }
    }
    
    /// 创建带参数的SQL查询
    pub fn sql_with_params(content: String, parameters: HashMap<String, QueryParameter>) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            query_type: QueryType::Sql,
            content,
            parameters,
            options: QueryOptions::default(),
            context: QueryContext::default(),
        }
    }
    
    /// 创建GraphQL查询
    pub fn graphql(content: String) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            query_type: QueryType::GraphQL,
            content,
            parameters: HashMap::new(),
            options: QueryOptions::default(),
            context: QueryContext::default(),
        }
    }
    
    /// 创建脚本查询
    pub fn script(content: String) -> Self {
        Self {
            id: uuid::Uuid::new_v4().to_string(),
            query_type: QueryType::Script,
            content,
            parameters: HashMap::new(),
            options: QueryOptions::default(),
            context: QueryContext::default(),
        }
    }
    
    /// 设置查询选项
    pub fn with_options(mut self, options: QueryOptions) -> Self {
        self.options = options;
        self
    }
    
    /// 设置查询上下文
    pub fn with_context(mut self, context: QueryContext) -> Self {
        self.context = context;
        self
    }
    
    /// 添加参数
    pub fn with_parameter(mut self, key: String, value: QueryParameter) -> Self {
        self.parameters.insert(key, value);
        self
    }
}

impl Default for QueryOptions {
    fn default() -> Self {
        Self {
            timeout_ms: Some(30000), // 30秒默认超时
            max_rows: Some(10000),   // 默认最大1万行
            streaming: false,
            cache_strategy: CacheStrategy::Auto,
            execution_mode: ExecutionMode::Sync,
        }
    }
}

impl Default for QueryContext {
    fn default() -> Self {
        Self {
            user_id: None,
            session_id: None,
            request_id: uuid::Uuid::new_v4().to_string(),
            datasource_id: None,
            permissions: Vec::new(),
            attributes: HashMap::new(),
        }
    }
}

impl QueryResult {
    /// 创建成功结果
    pub fn success(query_id: String, data: QueryResultData, stats: QueryStats) -> Self {
        Self {
            query_id,
            status: QueryStatus::Success,
            data,
            stats,
            error: None,
        }
    }
    
    /// 创建失败结果
    pub fn failed(query_id: String, error: String, stats: QueryStats) -> Self {
        Self {
            query_id,
            status: QueryStatus::Failed,
            data: QueryResultData::Empty,
            stats,
            error: Some(error),
        }
    }
    
    /// 检查是否成功
    pub fn is_success(&self) -> bool {
        matches!(self.status, QueryStatus::Success)
    }
    
    /// 获取Arrow数据
    pub fn as_arrow(&self) -> Option<&Vec<RecordBatch>> {
        match &self.data {
            QueryResultData::Arrow(batches) => Some(batches),
            _ => None,
        }
    }
    
    /// 获取JSON数据
    pub fn as_json(&self) -> Option<&serde_json::Value> {
        match &self.data {
            QueryResultData::Json(value) => Some(value),
            _ => None,
        }
    }
}

impl QueryStats {
    /// 创建新的统计信息
    pub fn new() -> Self {
        Self {
            start_time: chrono::Utc::now(),
            end_time: None,
            duration_ms: None,
            rows_processed: None,
            rows_returned: None,
            memory_used: None,
            cpu_time_ms: None,
            cache_hit: false,
        }
    }
    
    /// 标记查询完成
    pub fn finish(&mut self) {
        let end_time = chrono::Utc::now();
        self.end_time = Some(end_time);
        self.duration_ms = Some(
            (end_time - self.start_time).num_milliseconds() as u64
        );
    }
    
    /// 设置处理行数
    pub fn set_rows_processed(&mut self, count: usize) {
        self.rows_processed = Some(count);
    }
    
    /// 设置返回行数
    pub fn set_rows_returned(&mut self, count: usize) {
        self.rows_returned = Some(count);
    }
    
    /// 设置内存使用
    pub fn set_memory_used(&mut self, bytes: usize) {
        self.memory_used = Some(bytes);
    }
    
    /// 设置缓存命中
    pub fn set_cache_hit(&mut self, hit: bool) {
        self.cache_hit = hit;
    }
}

/// 查询执行器接口
#[async_trait::async_trait]
pub trait QueryExecutor: Send + Sync {
    /// 执行查询
    async fn execute(&self, query: Query) -> Result<QueryResult>;
    
    /// 取消查询
    async fn cancel(&self, query_id: &str) -> Result<()>;
    
    /// 获取查询状态
    async fn get_status(&self, query_id: &str) -> Result<QueryStatus>;
    
    /// 支持的查询类型
    fn supported_query_types(&self) -> Vec<QueryType>;
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_query_creation() {
        let query = Query::sql("SELECT * FROM users".to_string());
        assert_eq!(query.query_type, QueryType::Sql);
        assert_eq!(query.content, "SELECT * FROM users");
        assert!(query.parameters.is_empty());
    }
    
    #[test]
    fn test_query_with_parameters() {
        let mut params = HashMap::new();
        params.insert("id".to_string(), QueryParameter::Integer(123));
        params.insert("name".to_string(), QueryParameter::String("test".to_string()));
        
        let query = Query::sql_with_params(
            "SELECT * FROM users WHERE id = ? AND name = ?".to_string(),
            params,
        );
        
        assert_eq!(query.parameters.len(), 2);
        assert!(matches!(
            query.parameters.get("id"),
            Some(QueryParameter::Integer(123))
        ));
    }
    
    #[test]
    fn test_query_result() {
        let stats = QueryStats::new();
        let result = QueryResult::success(
            "test-query".to_string(),
            QueryResultData::Empty,
            stats,
        );
        
        assert!(result.is_success());
        assert_eq!(result.query_id, "test-query");
    }
    
    #[test]
    fn test_query_stats() {
        let mut stats = QueryStats::new();
        let start_time = stats.start_time;
        
        // 模拟一些处理时间
        std::thread::sleep(std::time::Duration::from_millis(1));
        
        stats.finish();
        stats.set_rows_processed(100);
        stats.set_rows_returned(50);
        stats.set_cache_hit(true);
        
        assert!(stats.end_time.is_some());
        assert!(stats.duration_ms.is_some());
        assert_eq!(stats.rows_processed, Some(100));
        assert_eq!(stats.rows_returned, Some(50));
        assert!(stats.cache_hit);
    }
}
